{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "11af8d6fd87cabcc",
   "metadata": {},
   "source": [
    "# Notebook Description\n",
    "\n",
    "Below is a complete training code for **AlterNet-LC**, including data preprocessing, model definition, training, validation, and testing steps.\n",
    "\n",
    "The notebook also includes comparative experiments on the **PneumoniaMNIST** dataset with three models: **DenseNet**, **GLCM+SVM**, and **LBP+RF**.\n",
    "\n",
    "Finally, the trained **AlterNet-LC** model is tested for generalization capability on a **Kaggle dataset**.\n",
    "\n",
    "#### **Warning:**\n",
    "The code in this notebook is for reference only. It has **not undergone strict data leakage prevention or logic optimization** and should **not be directly used in production environments**.\n",
    "\n",
    "Results may vary due to device differences. Default configurations are provided for reference; adjust and debug according to your actual setup.\n",
    "\n",
    "#### Authors：Li Jiawei；Chen Mingfang；Yao Zehan"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96e1a314b9c8830",
   "metadata": {},
   "source": [
    "MedMNIST provides a corresponding Python library for accessing its datasets.\n",
    "\n",
    "First, install `medmnist`:\n",
    "```bash\n",
    "pip install medmnist\n",
    "```\n",
    "\n",
    "Then download the dataset locally via Python (may require a VPN):\n",
    "```python\n",
    "import medmnist\n",
    "from medmnist import INFO\n",
    "import numpy as np\n",
    "import torch\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "data_transform = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=[.5], std=[.5])\n",
    "])\n",
    "\n",
    "data_flag = 'pneumoniamnist'\n",
    "info = INFO[data_flag]\n",
    "DataClass = getattr(medmnist, info['python_class'])\n",
    "\n",
    "train_dataset = DataClass(split='train', transform=data_transform, download=True, size=224, mmap_mode='r')\n",
    "```\n",
    "\n",
    "For generalization testing, the external test dataset can be downloaded via the Kaggle API. First, install `kagglehub`:\n",
    "```bash\n",
    "pip install kagglehub\n",
    "```\n",
    "\n",
    "The first time using `kagglehub` may require logging into your Kaggle account:\n",
    "```python\n",
    "import kagglehub\n",
    "kagglehub.login()\n",
    "```\n",
    "\n",
    "After logging in, download the dataset:\n",
    "```python\n",
    "import kagglehub\n",
    "\n",
    "# Download latest version\n",
    "path = kagglehub.dataset_download(\"paultimothymooney/chest-xray-pneumonia\")\n",
    "\n",
    "print(\"Path to dataset files:\", path)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7858e3a82a79814a",
   "metadata": {},
   "source": "# AlterNet-LC"
  },
  {
   "cell_type": "code",
   "id": "901246eacaf14f14",
   "metadata": {},
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "from torchvision import transforms, models\n",
    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
    "from sklearn.metrics import roc_curve, auc, precision_recall_curve, average_precision_score\n",
    "from sklearn.metrics import precision_score, recall_score, f1_score, confusion_matrix, classification_report\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from matplotlib.colors import LinearSegmentedColormap\n",
    "from tqdm import tqdm\n",
    "import os\n",
    "import datetime\n",
    "import time\n",
    "# Chinese Font Support\n",
    "plt.rcParams['font.sans-serif'] = ['SimHei']\n",
    "plt.rcParams['axes.unicode_minus'] = False\n",
    "# random number seed\n",
    "torch.manual_seed(101010)\n",
    "np.random.seed(101010)\n",
    "# Checking GPU availability\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print(f\"Device: {device}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "2db86b7debc8044a",
   "metadata": {},
   "source": [
    "# data path\n",
    "data_path = r\"pneumoniamnist_224.npz\"\n",
    "def ensure_dir(dir_path):\n",
    "    if not os.path.exists(dir_path):\n",
    "        os.makedirs(dir_path)\n",
    "        print(f\"make dir: {dir_path}\")\n",
    "\n",
    "# make results_dir\n",
    "results_dir = \"evaluation_results\"\n",
    "ensure_dir(results_dir)\n",
    "img_dir = os.path.join(results_dir, \"images\")\n",
    "ensure_dir(img_dir)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "933595c66fa490eb",
   "metadata": {},
   "source": "## 1.WindowAttention"
  },
  {
   "cell_type": "code",
   "id": "98b37b2f1e5787fb",
   "metadata": {},
   "source": [
    "# WindowAttention\n",
    "class WindowAttention:\n",
    "    @staticmethod\n",
    "    def partition(x, window_size):\n",
    "        \"\"\"\n",
    "        Partition the input tensor into non-overlapping windows.\n",
    "        Args:\n",
    "            x: input tensor (B, C, H, W)\n",
    "            window_size: window size\n",
    "        Returns:\n",
    "            windows: windowed tensor (B*num_windows, C, window_size, window_size)\n",
    "            padded_shape: padded shape (Hp, Wp)\n",
    "        \"\"\"\n",
    "        B, C, H, W = x.shape\n",
    "\n",
    "        pad_h = (window_size - H % window_size) % window_size\n",
    "        pad_w = (window_size - W % window_size) % window_size\n",
    "\n",
    "        if pad_h > 0 or pad_w > 0:\n",
    "            x = F.pad(x, (0, pad_w, 0, pad_h))\n",
    "\n",
    "        Hp, Wp = H + pad_h, W + pad_w\n",
    "        num_win_h, num_win_w = Hp // window_size, Wp // window_size\n",
    "\n",
    "        x = x.view(B, C, num_win_h, window_size, num_win_w, window_size)\n",
    "\n",
    "        windows = x.permute(0, 2, 4, 1, 3, 5).contiguous().view(-1, C, window_size, window_size)\n",
    "\n",
    "        return windows, (Hp, Wp)\n",
    "\n",
    "    @staticmethod\n",
    "    def reverse(windows, window_size, H, W, Hp, Wp):\n",
    "        \"\"\"\n",
    "        Merge windows back into the complete feature map\n",
    "        Args:\n",
    "            windows: windowed tensor (B*num_windows, C, window_size, window_size)\n",
    "            window_size: window size\n",
    "            H, W: original height and width\n",
    "            Hp, Wp: padded height and width\n",
    "        Returns:\n",
    "            merged tensor (B, C, H, W)\n",
    "        \"\"\"\n",
    "        B_win, C, _, _ = windows.shape\n",
    "        num_win_h, num_win_w = Hp // window_size, Wp // window_size\n",
    "        B = B_win // (num_win_h * num_win_w)\n",
    "\n",
    "        x = windows.view(B, num_win_h, num_win_w, C, window_size, window_size)\n",
    "\n",
    "        x = x.permute(0, 3, 1, 4, 2, 5).contiguous()\n",
    "\n",
    "        x = x.view(B, C, Hp, Wp)\n",
    "\n",
    "        if Hp > H or Wp > W:\n",
    "            x = x[:, :, :H, :W]\n",
    "\n",
    "        return x"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "e31b7d86e4397d67",
   "metadata": {},
   "source": "## 2.Multiple Self-attention Blocks"
  },
  {
   "cell_type": "code",
   "id": "2ff90e4fce9c0cae",
   "metadata": {},
   "source": [
    "class MSABlock(nn.Module):\n",
    "    \"\"\"Multiple self-attention blocks, include Window Attention, LayerNorm and Residual Connection\"\"\"\n",
    "\n",
    "    def __init__(self, dim, num_heads, head_dim, window_size=7, mlp_ratio=4., dropout_rate=0.1):\n",
    "        super().__init__()\n",
    "        self.num_heads = num_heads\n",
    "        self.head_dim = head_dim\n",
    "        self.window_size = window_size\n",
    "        self.mlp_ratio = mlp_ratio\n",
    "        self.dropout_rate = dropout_rate\n",
    "        self.total_head_dim = num_heads * head_dim\n",
    "\n",
    "        # attention layer\n",
    "        self.qkv = nn.Linear(dim, self.total_head_dim * 3, bias=False)\n",
    "        self.proj = nn.Linear(self.total_head_dim, dim)\n",
    "        self.proj_drop = nn.Dropout(dropout_rate)\n",
    "\n",
    "        # MLP layer\n",
    "        mlp_hidden_dim = int(dim * mlp_ratio)\n",
    "        self.mlp = nn.Sequential(\n",
    "            nn.Linear(dim, mlp_hidden_dim),\n",
    "            nn.GELU(),\n",
    "            nn.Dropout(dropout_rate),\n",
    "            nn.Linear(mlp_hidden_dim, dim),\n",
    "            nn.Dropout(dropout_rate)\n",
    "        )\n",
    "\n",
    "        # normalization layer\n",
    "        self.norm1 = nn.LayerNorm(dim, eps=1e-6)\n",
    "        self.norm2 = nn.LayerNorm(dim, eps=1e-6)\n",
    "\n",
    "        # Weight Initialization\n",
    "        nn.init.xavier_uniform_(self.qkv.weight)\n",
    "        nn.init.xavier_uniform_(self.proj.weight)\n",
    "\n",
    "    def forward(self, x):\n",
    "        B, C, H, W = x.shape\n",
    "        shortcut1 = x\n",
    "\n",
    "        # normalization\n",
    "        x = x.permute(0, 2, 3, 1)  # B, H, W, C\n",
    "        x = self.norm1(x)\n",
    "\n",
    "        # Window segmentation\n",
    "        x_windows, (Hp, Wp) = WindowAttention.partition(x.permute(0, 3, 1, 2), self.window_size)\n",
    "        x_windows = x_windows.permute(0, 2, 3, 1)\n",
    "        win_seq_len = self.window_size * self.window_size\n",
    "        x_windows_seq = x_windows.reshape(-1, win_seq_len, C)\n",
    "\n",
    "        # generate QKV\n",
    "        qkv = self.qkv(x_windows_seq)\n",
    "        qkv = qkv.reshape(-1, win_seq_len, 3, self.num_heads, self.head_dim)\n",
    "        q, k, v = qkv.unbind(dim=2)\n",
    "\n",
    "        # attention calculation\n",
    "        q = q.transpose(1, 2)\n",
    "        k = k.transpose(1, 2)\n",
    "        v = v.transpose(1, 2)\n",
    "\n",
    "        scale = self.head_dim ** -0.5\n",
    "        attn = (q @ k.transpose(-2, -1)) * scale\n",
    "        attn = attn.softmax(dim=-1)\n",
    "\n",
    "        # weighted sum\n",
    "        attn_output = (attn @ v)\n",
    "        attn_output = attn_output.transpose(1, 2)\n",
    "        attn_output = attn_output.reshape(-1, win_seq_len, self.total_head_dim)\n",
    "\n",
    "        # projection\n",
    "        attn_output = self.proj(attn_output)\n",
    "        attn_output = self.proj_drop(attn_output)\n",
    "\n",
    "        # Restore window format\n",
    "        attn_output = attn_output.reshape(-1, self.window_size, self.window_size, C)\n",
    "        attn_output = attn_output.permute(0, 3, 1, 2)\n",
    "\n",
    "        # merge window\n",
    "        attn_merged = WindowAttention.reverse(attn_output, self.window_size, H, W, Hp, Wp)\n",
    "\n",
    "        # first residual link\n",
    "        x = shortcut1 + attn_merged\n",
    "\n",
    "        # MLP\n",
    "        shortcut2 = x\n",
    "        x = x.permute(0, 2, 3, 1)\n",
    "        x = self.norm2(x)\n",
    "        x = self.mlp(x)\n",
    "        x = x.permute(0, 3, 1, 2)\n",
    "\n",
    "        # second residual link\n",
    "        x = shortcut2 + x\n",
    "\n",
    "        return x"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "c9788b9d16f3f30e",
   "metadata": {},
   "source": "## 3.Contrast Self-attention Blocks"
  },
  {
   "cell_type": "code",
   "id": "3d41a93fedfbc851",
   "metadata": {},
   "source": [
    "class ContrastiveMSABlock(MSABlock):\n",
    "    \"\"\"Multiple Self-Attention Blocks with Contrastive Learning Headers\"\"\"\n",
    "\n",
    "    def __init__(self, dim, num_heads, head_dim, window_size=7, mlp_ratio=4., dropout_rate=0.1,\n",
    "                 proj_dim=128, temperature=0.07):\n",
    "        super().__init__(dim, num_heads, head_dim, window_size, mlp_ratio, dropout_rate)\n",
    "\n",
    "        # Compare and contrast learning parameters\n",
    "        self.temperature = temperature\n",
    "\n",
    "        # Comparative Learning Projection Heads\n",
    "        self.contrast_projection = nn.Sequential(\n",
    "            nn.Linear(dim, dim // 2),\n",
    "            nn.LayerNorm(dim // 2),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Linear(dim // 2, proj_dim)\n",
    "        )\n",
    "\n",
    "        # Initialize projection header weights\n",
    "        for m in self.contrast_projection.modules():\n",
    "            if isinstance(m, nn.Linear):\n",
    "                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n",
    "                if m.bias is not None:\n",
    "                    nn.init.constant_(m.bias, 0)\n",
    "\n",
    "    def contrastive_loss(self, features, labels):\n",
    "\n",
    "        # Normalize the features\n",
    "        features = F.normalize(features, dim=1)\n",
    "\n",
    "        # Calculate the cosine similarity between samples\n",
    "        similarity_matrix = torch.matmul(features, features.T) / self.temperature\n",
    "\n",
    "        # Creating a label similarity matrix\n",
    "        labels = labels.view(-1, 1)\n",
    "        mask_same_class = torch.eq(labels, labels.T).float()\n",
    "\n",
    "        # Creating a self-mask\n",
    "        mask_self = torch.eye(mask_same_class.shape[0], device=mask_same_class.device)\n",
    "        mask_other = 1 - mask_self\n",
    "\n",
    "        # positive sample-pair mask (computing)\n",
    "        mask_positives = mask_same_class * mask_other\n",
    "\n",
    "        # Check if there are enough positive sample pairs before calculating the loss\n",
    "        positive_per_sample = mask_positives.sum(1)\n",
    "        valid_samples = positive_per_sample > 0\n",
    "\n",
    "        if valid_samples.sum() == 0:\n",
    "            return torch.tensor(0.0, device=features.device)\n",
    "\n",
    "        # Calculation of comparative losses\n",
    "        exp_similarity = torch.exp(similarity_matrix) * mask_other\n",
    "\n",
    "        # Calculate log-sum-exp for each sample\n",
    "        log_prob_denominator = torch.log(exp_similarity.sum(1, keepdim=True) + 1e-12)\n",
    "\n",
    "        # Calculate the log probability of a positive sample pair\n",
    "        log_prob_positives = similarity_matrix - log_prob_denominator\n",
    "\n",
    "        # Calculate the average log probability of a positive sample pair for each sample\n",
    "        mean_log_prob_positives = (mask_positives * log_prob_positives).sum(1) / (positive_per_sample + 1e-12)\n",
    "\n",
    "        # Calculation of ultimate loss\n",
    "        mean_log_prob_positives = mean_log_prob_positives[valid_samples]\n",
    "\n",
    "        # Final loss\n",
    "        contrastive_loss = -mean_log_prob_positives.mean()\n",
    "\n",
    "        return contrastive_loss\n",
    "\n",
    "    def forward(self, x, labels=None):\n",
    "        \"\"\"Forward Propagation\n",
    "\n",
    "        Args:\n",
    "            x: Input features [B, C, H, W]\n",
    "            labels: Optional, class labels [B]\n",
    "\n",
    "        Returns:\n",
    "            x: Processed features\n",
    "            contrast_loss: Contrastive loss value\n",
    "        \"\"\"\n",
    "        # Call forward propagation of the parent class to get the basic output\n",
    "        x = super().forward(x)\n",
    "\n",
    "        # If no labels are provided, return directly to the feature\n",
    "        if labels is None:\n",
    "            return x, 0.0\n",
    "\n",
    "        # Global average pooling extracts a feature representation for each sample\n",
    "        batch_size = x.shape[0]\n",
    "        pooled_features = F.adaptive_avg_pool2d(x, (1, 1)).view(batch_size, -1)\n",
    "\n",
    "        # Mapping to Contrast Learning Space using a Projection Head\n",
    "        projected_features = self.contrast_projection(pooled_features)\n",
    "\n",
    "        # Calculation of comparative losses\n",
    "        contrast_loss = self.contrastive_loss(projected_features, labels)\n",
    "\n",
    "        return x, contrast_loss"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "328c1f5cf40b8e60",
   "metadata": {},
   "source": "## 4.Pre-activated Residual Block"
  },
  {
   "cell_type": "code",
   "id": "2e235e4fdbc82ad5",
   "metadata": {},
   "source": [
    "class PreActResidualBlock(nn.Module):\n",
    "    \"\"\"Pre-activated residual block\"\"\"\n",
    "\n",
    "    def __init__(self, in_channels, out_channels, stride=1, use_projection=False):\n",
    "        super().__init__()\n",
    "        self.use_projection = use_projection\n",
    "\n",
    "        # Projection Shortcut\n",
    "        if self.use_projection:\n",
    "            self.proj_conv = nn.Conv2d(\n",
    "                in_channels, out_channels, kernel_size=1,\n",
    "                stride=stride, padding=0, bias=False\n",
    "            )\n",
    "\n",
    "        # primary path\n",
    "        self.bn1 = nn.BatchNorm2d(in_channels)\n",
    "        self.conv1 = nn.Conv2d(\n",
    "            in_channels, out_channels, kernel_size=3,\n",
    "            stride=stride, padding=1, bias=False\n",
    "        )\n",
    "        self.bn2 = nn.BatchNorm2d(out_channels)\n",
    "        self.conv2 = nn.Conv2d(\n",
    "            out_channels, out_channels, kernel_size=3,\n",
    "            stride=1, padding=1, bias=False\n",
    "        )\n",
    "\n",
    "        # Weight Initialisation\n",
    "        nn.init.kaiming_normal_(self.conv1.weight, mode='fan_out', nonlinearity='relu')\n",
    "        nn.init.kaiming_normal_(self.conv2.weight, mode='fan_out', nonlinearity='relu')\n",
    "        if self.use_projection:\n",
    "            nn.init.kaiming_normal_(self.proj_conv.weight, mode='fan_out', nonlinearity='relu')\n",
    "\n",
    "    def forward(self, x):\n",
    "        shortcut = x\n",
    "\n",
    "        # pre-activation\n",
    "        x = self.bn1(x)\n",
    "        x = F.relu(x)\n",
    "\n",
    "        # Projection Shortcut\n",
    "        if self.use_projection:\n",
    "            shortcut = self.proj_conv(x)\n",
    "\n",
    "        # The first convolution\n",
    "        x = self.conv1(x)\n",
    "        x = self.bn2(x)\n",
    "        x = F.relu(x)\n",
    "\n",
    "        # The second convolution\n",
    "        x = self.conv2(x)\n",
    "\n",
    "        # Shortcut Connection\n",
    "        return shortcut + x"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "e223685a70f61df8",
   "metadata": {},
   "source": "## 5.AlterNet-LC Model Architecture"
  },
  {
   "cell_type": "code",
   "id": "5e43eb19f020283e",
   "metadata": {},
   "source": [
    "class AlterNet_LC(nn.Module):\n",
    "    \"\"\"AlterNet-LC model\"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        in_channels=3,\n",
    "        num_classes=2,\n",
    "        blocks_per_stage=[3, 4, 6, 3],\n",
    "        initial_filters=128,\n",
    "        head_counts=[6, 12, 24],\n",
    "        head_dim=64,\n",
    "        window_size=7,\n",
    "        dropout_rate=0.3\n",
    "    ):\n",
    "        super().__init__()\n",
    "\n",
    "        # Stem\n",
    "        self.stem = nn.Sequential(\n",
    "            nn.Conv2d(in_channels, initial_filters, kernel_size=7, stride=2, padding=3, bias=False),\n",
    "            nn.BatchNorm2d(initial_filters),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n",
    "        )\n",
    "\n",
    "        # networking phase\n",
    "        self.stages = nn.ModuleList()\n",
    "        current_filters = initial_filters\n",
    "\n",
    "        # Stage 1 (pure CNN)\n",
    "        stage1 = nn.Sequential()\n",
    "        stage1.add_module('block0', PreActResidualBlock(current_filters, current_filters, use_projection=True))\n",
    "        for i in range(1, blocks_per_stage[0]):\n",
    "            stage1.add_module(f'block{i}', PreActResidualBlock(current_filters, current_filters))\n",
    "        self.stages.append(stage1)\n",
    "\n",
    "        # Stages 2-4 (CNN + MSA with ContrastiveHead)\n",
    "        self.contrastive_blocks = nn.ModuleList()\n",
    "\n",
    "        for stage_idx, num_blocks in enumerate(blocks_per_stage[1:]):\n",
    "            stage_num = stage_idx + 2\n",
    "            next_filters = current_filters * 2\n",
    "            num_heads = head_counts[stage_idx]\n",
    "\n",
    "            stage = nn.Sequential()\n",
    "\n",
    "            # First block of the phase (downsampling)\n",
    "            stage.add_module(\n",
    "                'block0',\n",
    "                PreActResidualBlock(current_filters, next_filters, stride=2, use_projection=True)\n",
    "            )\n",
    "\n",
    "            # Intermediate CNN blocks\n",
    "            for i in range(1, num_blocks - 1):\n",
    "                stage.add_module(f'block{i}', PreActResidualBlock(next_filters, next_filters))\n",
    "\n",
    "            # The last block is MSA with comparative learning\n",
    "            contrastive_block = ContrastiveMSABlock(\n",
    "                dim=next_filters,\n",
    "                num_heads=num_heads,\n",
    "                head_dim=head_dim,\n",
    "                window_size=window_size,\n",
    "                temperature=0.07\n",
    "            )\n",
    "            stage.add_module(f'block{num_blocks-1}', contrastive_block)\n",
    "            self.contrastive_blocks.append(contrastive_block)\n",
    "\n",
    "            self.stages.append(stage)\n",
    "            current_filters = next_filters\n",
    "\n",
    "        # Classification Head\n",
    "        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n",
    "        self.dropout = nn.Dropout(dropout_rate)\n",
    "        self.fc = nn.Linear(current_filters, num_classes)\n",
    "\n",
    "        # Weight Initialisation\n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.Conv2d):\n",
    "                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n",
    "            elif isinstance(m, nn.BatchNorm2d):\n",
    "                nn.init.constant_(m.weight, 1)\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "\n",
    "    def forward(self, x, labels=None):\n",
    "        x = self.stem(x)\n",
    "\n",
    "        # Initialising the sum of comparison losses\n",
    "        contrastive_loss_sum = 0.0\n",
    "\n",
    "        # Stage 1\n",
    "        x = self.stages[0](x)\n",
    "\n",
    "        # Stage 2-4\n",
    "        for i, stage in enumerate(self.stages[1:], 1):\n",
    "            for name, block in stage.named_children():\n",
    "                if name != f'block{len(stage)-1}':\n",
    "                    x = block(x)\n",
    "                else:\n",
    "                    x, stage_contrast_loss = block(x, labels)\n",
    "                    contrastive_loss_sum += stage_contrast_loss\n",
    "\n",
    "        # Classification Head\n",
    "        x = self.avgpool(x)\n",
    "        x = torch.flatten(x, 1)\n",
    "        x = self.dropout(x)\n",
    "        x = self.fc(x)\n",
    "\n",
    "        # If it is a train pattern and labels are provided, return the classification output and comparison loss\n",
    "        if self.training and labels is not None:\n",
    "            return x, contrastive_loss_sum\n",
    "        else:\n",
    "            return x"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "80f3aa7e96b8f082",
   "metadata": {},
   "source": "## 6.Transformation of Pneumonia Imaging Dataset"
  },
  {
   "cell_type": "code",
   "id": "81ccab3412e21898",
   "metadata": {},
   "source": [
    "class MedicalDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, images, labels, transform=None):\n",
    "        self.images = images\n",
    "        self.labels = labels\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.images)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        img = self.images[idx]  # original shape: (C, H, W)\n",
    "        label = self.labels[idx]\n",
    "\n",
    "        if self.transform:\n",
    "            # Convert to PIL Image\n",
    "            img = self.transform(img.numpy().transpose(1,2,0))\n",
    "\n",
    "        return img.float(), label"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "851d2311d652f833",
   "metadata": {},
   "source": "## 7.Load and Preprocess Dataset Function"
  },
  {
   "cell_type": "code",
   "id": "ba52d6d33c545faf",
   "metadata": {},
   "source": [
    "def load_and_preprocess_data(data_path, batch_size=64):\n",
    "\n",
    "    # Extract data from npz file\n",
    "    data = np.load(data_path)\n",
    "    train_images = data['train_images']\n",
    "    train_labels = data['train_labels'].squeeze()\n",
    "    test_images = data['test_images']\n",
    "    test_labels = data['test_labels'].squeeze()\n",
    "    val_images = data['val_images']\n",
    "    val_labels = data['val_labels'].squeeze()\n",
    "\n",
    "    # Check data distribution\n",
    "    print(\"Original label values:\", np.unique(data['train_labels']))\n",
    "    print(\"Training labels distribution:\", np.unique(train_labels, return_counts=True))\n",
    "    print(\"Validation labels distribution:\", np.unique(val_labels, return_counts=True))\n",
    "    print(\"Test labels distribution:\", np.unique(test_labels, return_counts=True))\n",
    "\n",
    "    # Ensure images have 3 channels (RGB)\n",
    "    def ensure_3channel(images):\n",
    "        if images.ndim == 3:\n",
    "            images = np.expand_dims(images, axis=-1)\n",
    "        if images.shape[-1] == 1:\n",
    "            images = np.repeat(images, 3, axis=-1)\n",
    "        return images\n",
    "\n",
    "    train_images = ensure_3channel(train_images)\n",
    "    val_images = ensure_3channel(val_images)\n",
    "    test_images = ensure_3channel(test_images)\n",
    "\n",
    "    # Normalize and convert to PyTorch format (N, C, H, W)\n",
    "    def preprocess_images(images):\n",
    "        images = np.transpose(images, (0, 3, 1, 2))\n",
    "        return images\n",
    "\n",
    "    train_images = preprocess_images(train_images)\n",
    "    val_images = preprocess_images(val_images)\n",
    "    test_images = preprocess_images(test_images)\n",
    "\n",
    "    # Data augmentation transforms\n",
    "    train_transform = transforms.Compose([\n",
    "        transforms.ToPILImage(),\n",
    "        transforms.RandomHorizontalFlip(p=0.5),\n",
    "        transforms.RandomRotation(15),\n",
    "        transforms.ColorJitter(brightness=0.2, contrast=0.2),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(mean=[0.5], std=[0.5])\n",
    "    ])\n",
    "\n",
    "    basic_transform = transforms.Compose([\n",
    "        transforms.ToPILImage(),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(mean=[0.5], std=[0.5])\n",
    "    ])\n",
    "\n",
    "    # Create PyTorch datasets\n",
    "    train_dataset = MedicalDataset(\n",
    "        torch.tensor(train_images, dtype=torch.uint8),\n",
    "        torch.tensor(train_labels, dtype=torch.long),\n",
    "        transform=train_transform\n",
    "    )\n",
    "\n",
    "    val_dataset = MedicalDataset(\n",
    "        torch.tensor(val_images, dtype=torch.uint8),\n",
    "        torch.tensor(val_labels, dtype=torch.long),\n",
    "        transform=basic_transform\n",
    "    )\n",
    "\n",
    "    test_dataset = MedicalDataset(\n",
    "        torch.tensor(test_images, dtype=torch.uint8),\n",
    "        torch.tensor(test_labels, dtype=torch.long),\n",
    "        transform=basic_transform\n",
    "    )\n",
    "\n",
    "    # Create data loaders\n",
    "    train_loader = DataLoader(\n",
    "        train_dataset,\n",
    "        batch_size=batch_size,\n",
    "        shuffle=True,\n",
    "        pin_memory=True,\n",
    "    )\n",
    "    val_loader = DataLoader(\n",
    "        val_dataset,\n",
    "        batch_size=batch_size,\n",
    "        shuffle=False,\n",
    "        pin_memory=True,\n",
    "    )\n",
    "    test_loader = DataLoader(\n",
    "        test_dataset,\n",
    "        batch_size=batch_size,\n",
    "        shuffle=False,\n",
    "        pin_memory=True,\n",
    "    )\n",
    "\n",
    "    return train_loader, val_loader, test_loader"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "ffffd3f945a30278",
   "metadata": {},
   "source": "## 8.LMF Loss"
  },
  {
   "cell_type": "code",
   "id": "dac6f261afc5a616",
   "metadata": {},
   "source": [
    "class LMFLoss(nn.Module):\n",
    "    def __init__(self, class_weights, gamma=2.6, margin=0.7, alpha=0.27):\n",
    "        super().__init__()\n",
    "        self.gamma = gamma\n",
    "        self.margin = margin\n",
    "        self.alpha = alpha\n",
    "        self.class_weights = class_weights\n",
    "\n",
    "    def forward(self, inputs, targets):\n",
    "        ce_loss = F.cross_entropy(inputs, targets, weight=self.class_weights)\n",
    "        pt = torch.exp(-ce_loss)\n",
    "        focal_loss = (1 - pt)**self.gamma * ce_loss\n",
    "\n",
    "        # Margin loss\n",
    "        targets_onehot = F.one_hot(targets, num_classes=2).float()\n",
    "        probs = inputs.softmax(dim=1)\n",
    "\n",
    "        # Applying a boundary penalty for negative samples\n",
    "        neg_margin = self.margin * 1.15\n",
    "        pos_margin = self.margin\n",
    "\n",
    "        # Choose different margins depending on the target category\n",
    "        margins = torch.zeros_like(targets_onehot)\n",
    "        margins[:, 0] = neg_margin\n",
    "        margins[:, 1] = pos_margin\n",
    "\n",
    "        # Select the actual margin value to be used\n",
    "        selected_margins = (margins * targets_onehot).sum(1, keepdim=True)\n",
    "        margins = selected_margins * (1 - torch.abs(targets_onehot - probs))\n",
    "\n",
    "        margin_loss = F.relu(1 - (probs * (2*targets_onehot-1) - margins)).mean()\n",
    "\n",
    "        # Adding extra weight for misclassified negative samples\n",
    "        neg_samples = (targets == 0)\n",
    "        misclassified_neg = neg_samples & (probs[:, 1] > 0.25)\n",
    "\n",
    "        sample_weights = torch.ones_like(targets, dtype=torch.float32)\n",
    "        sample_weights[misclassified_neg] = 1.1\n",
    "\n",
    "        weighted_loss = (focal_loss + self.alpha*margin_loss) * sample_weights\n",
    "\n",
    "        return weighted_loss.mean()"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "a15c793676dc36a4",
   "metadata": {},
   "source": "## 9.Train_model Function"
  },
  {
   "cell_type": "code",
   "id": "96ef88580fac4511",
   "metadata": {},
   "source": [
    "def train_model(model, train_loader, val_loader, epochs=100, lr=5e-5, weight_decay=1.6e-4):\n",
    "    \"\"\"Train and validate a PyTorch model with contrastive learning\"\"\"\n",
    "\n",
    "    # Optimizer\n",
    "    optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)\n",
    "\n",
    "    # Learning rate scheduler\n",
    "    scheduler = ReduceLROnPlateau(\n",
    "        optimizer, mode='min', factor=0.5, patience=5, min_lr=1e-6, verbose=True\n",
    "    )\n",
    "\n",
    "    # Loss function\n",
    "    criterion = LMFLoss(\n",
    "        class_weights=torch.tensor([1.3, 1.25]).to(device),\n",
    "        gamma=2.6,\n",
    "        margin=0.7,\n",
    "        alpha=0.27\n",
    "    )\n",
    "\n",
    "    # Contrastive learning weight\n",
    "    contrastive_weight = 0.45\n",
    "\n",
    "    # Track best model\n",
    "    best_val_loss = float('inf')\n",
    "    best_model_state = None\n",
    "    patience_counter = 0\n",
    "    max_patience = 15\n",
    "\n",
    "    # Record training history\n",
    "    history = {\n",
    "        'train_loss': [], 'train_acc': [], 'train_auc': [],\n",
    "        'val_loss': [], 'val_acc': [], 'val_auc': []\n",
    "    }\n",
    "\n",
    "    # Training loop\n",
    "    for epoch in range(epochs):\n",
    "        # Training phase\n",
    "        model.train()\n",
    "        train_loss, train_correct, train_total = 0.0, 0, 0\n",
    "        train_outputs_all, train_labels_all = [], []\n",
    "\n",
    "        for inputs, labels in tqdm(train_loader, desc=f\"Epoch {epoch+1}/{epochs} [Train]\"):\n",
    "            inputs = inputs.to(device)\n",
    "            labels = labels.to(device)\n",
    "\n",
    "            # Forward pass\n",
    "            optimizer.zero_grad()\n",
    "            outputs, contrastive_loss = model(inputs, labels)\n",
    "\n",
    "            # Calculate LMF loss\n",
    "            lmf_loss = criterion(outputs, labels)\n",
    "\n",
    "            # Total loss = LMF loss + contrastive learning weight * contrastive loss\n",
    "            total_loss = lmf_loss + contrastive_weight * contrastive_loss\n",
    "\n",
    "            # Backward pass\n",
    "            total_loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            # Statistics collection\n",
    "            train_loss += lmf_loss.item() * inputs.size(0)\n",
    "            _, preds = torch.max(outputs, 1)\n",
    "            train_correct += torch.sum(preds == labels).item()\n",
    "            train_total += inputs.size(0)\n",
    "            train_outputs_all.append(F.softmax(outputs, dim=1).detach().cpu().numpy())\n",
    "            train_labels_all.append(labels.cpu().numpy())\n",
    "\n",
    "        # Calculate epoch-level training metrics\n",
    "        train_loss = train_loss / train_total\n",
    "        train_acc = train_correct / train_total\n",
    "\n",
    "        # Merge all batch outputs and labels\n",
    "        train_outputs_all = np.vstack(train_outputs_all)\n",
    "        train_labels_all = np.concatenate(train_labels_all)\n",
    "\n",
    "        # Calculate training AUC metrics\n",
    "        train_fpr, train_tpr, _ = roc_curve(\n",
    "            (train_labels_all == 1).astype(int),\n",
    "            train_outputs_all[:, 1]\n",
    "        )\n",
    "        train_auc_score = auc(train_fpr, train_tpr)\n",
    "\n",
    "        # Validation phase\n",
    "        model.eval()\n",
    "        val_loss, val_correct, val_total = 0.0, 0, 0\n",
    "        val_outputs_all, val_labels_all = [], []\n",
    "\n",
    "        with torch.no_grad():\n",
    "            for inputs, labels in tqdm(val_loader, desc=f\"Epoch {epoch+1}/{epochs} [Val]\"):\n",
    "                inputs = inputs.to(device)\n",
    "                labels = labels.to(device)\n",
    "                outputs = model(inputs)\n",
    "\n",
    "                # Calculate LMF loss\n",
    "                lmf_loss = criterion(outputs, labels)\n",
    "\n",
    "                # Only use classification loss in validation phase\n",
    "                val_loss += lmf_loss.item() * inputs.size(0)\n",
    "\n",
    "                _, preds = torch.max(outputs, 1)\n",
    "                val_correct += torch.sum(preds == labels).item()\n",
    "                val_total += inputs.size(0)\n",
    "                val_outputs_all.append(F.softmax(outputs, dim=1).detach().cpu().numpy())\n",
    "                val_labels_all.append(labels.cpu().numpy())\n",
    "\n",
    "        # Calculate epoch-level validation metrics\n",
    "        val_loss = val_loss / val_total\n",
    "        val_acc = val_correct / val_total\n",
    "\n",
    "        # Merge all batch outputs and labels\n",
    "        val_outputs_all = np.vstack(val_outputs_all)\n",
    "        val_labels_all = np.concatenate(val_labels_all)\n",
    "\n",
    "        # Calculate validation AUC\n",
    "        val_fpr, val_tpr, _ = roc_curve(\n",
    "            (val_labels_all == 1).astype(int),\n",
    "            val_outputs_all[:, 1]\n",
    "        )\n",
    "        val_auc_score = auc(val_fpr, val_tpr)\n",
    "\n",
    "        # Update learning rate\n",
    "        scheduler.step(val_loss)\n",
    "\n",
    "        # Update history records\n",
    "        history['train_loss'].append(train_loss)\n",
    "        history['train_acc'].append(train_acc)\n",
    "        history['train_auc'].append(train_auc_score)\n",
    "        history['val_loss'].append(val_loss)\n",
    "        history['val_acc'].append(val_acc)\n",
    "        history['val_auc'].append(val_auc_score)\n",
    "\n",
    "        # Output current epoch results\n",
    "        print(f\"Epoch {epoch+1}/{epochs}:\")\n",
    "        print(f\"  Train - Loss: {train_loss:.4f}, Acc: {train_acc:.4f}, AUC: {train_auc_score:.4f}\")\n",
    "        print(f\"  Val   - Loss: {val_loss:.4f}, Acc: {val_acc:.4f}, AUC: {val_auc_score:.4f}\")\n",
    "\n",
    "        # Save best model\n",
    "        if val_loss < best_val_loss:\n",
    "            best_val_loss = val_loss\n",
    "            best_model_state = model.state_dict().copy()\n",
    "            patience_counter = 0\n",
    "            print(f\"  New best model saved! (Val Loss: {val_loss:.4f})\")\n",
    "\n",
    "            # Save model file\n",
    "            torch.save(best_model_state, f'alternet_contrastive_{datetime.datetime.now().strftime(\"%Y%m%d\")}.pth')\n",
    "        else:\n",
    "            patience_counter += 1\n",
    "            print(f\"  No improvement for {patience_counter} epochs\")\n",
    "\n",
    "        # Early stopping\n",
    "        if patience_counter >= max_patience:\n",
    "            print(f\"Early stopping triggered after {epoch+1} epochs\")\n",
    "            break\n",
    "\n",
    "    # Load best model\n",
    "    if best_model_state is not None:\n",
    "        model.load_state_dict(best_model_state)\n",
    "\n",
    "    return model, history\n"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "c74a22c395a8f7b2",
   "metadata": {},
   "source": "## 10.Training Curve Plotting Function"
  },
  {
   "cell_type": "code",
   "id": "ab6caa8ae266ac0a",
   "metadata": {},
   "source": [
    "def plot_training(history):\n",
    "    plt.figure(figsize=(16, 10))\n",
    "\n",
    "    # Accuracy curve\n",
    "    plt.subplot(2, 2, 1)\n",
    "    plt.plot(history['train_acc'], label='Training Accuracy')\n",
    "    plt.plot(history['val_acc'], label='Validation Accuracy')\n",
    "    plt.title('Accuracy vs. Epochs')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Accuracy')\n",
    "    plt.legend()\n",
    "\n",
    "    # Loss curve\n",
    "    plt.subplot(2, 2, 2)\n",
    "    plt.plot(history['train_loss'], label='Training Loss')\n",
    "    plt.plot(history['val_loss'], label='Validation Loss')\n",
    "    plt.title('Loss vs. Epochs')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Loss')\n",
    "    plt.legend()\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "18aaf0cf9ca9320c",
   "metadata": {},
   "source": "## 11.Model Evaluation & Result Saving"
  },
  {
   "cell_type": "code",
   "id": "c8d6ed8d820540dc",
   "metadata": {},
   "source": [
    "def evaluate_model(model, test_loader, save_results=False):\n",
    "\n",
    "    # Get timestamp for result identification\n",
    "    timestamp = datetime.datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
    "\n",
    "    model.eval()\n",
    "    test_loss = 0.0\n",
    "    test_correct = 0\n",
    "    test_total = 0\n",
    "    all_preds = []\n",
    "    all_labels = []\n",
    "    all_outputs = []\n",
    "    all_inputs = []\n",
    "\n",
    "    # loss function\n",
    "    criterion = LMFLoss(\n",
    "        class_weights=torch.tensor([1.3, 1.25]).to(device),\n",
    "        gamma=2.6,\n",
    "        margin=0.7,\n",
    "        alpha=0.27\n",
    "    )\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for inputs, labels in tqdm(test_loader, desc=\"Evaluating Model\"):\n",
    "            inputs = inputs.to(device)\n",
    "            labels = labels.to(device)\n",
    "\n",
    "            # No need to calculate contrastive loss during evaluation\n",
    "            outputs = model(inputs)\n",
    "\n",
    "            # Handle tuple output (output, contrastive_loss) if present\n",
    "            if isinstance(outputs, tuple):\n",
    "                outputs = outputs[0]\n",
    "\n",
    "            loss = criterion(outputs, labels)\n",
    "\n",
    "            # Custom threshold prediction\n",
    "            probs = F.softmax(outputs, dim=1)[:, 1]\n",
    "            preds = (probs >= 0.985).long()\n",
    "\n",
    "            test_loss += loss.item() * inputs.size(0)\n",
    "            test_correct += torch.sum(preds == labels).item()\n",
    "            test_total += inputs.size(0)\n",
    "\n",
    "            all_preds.extend(preds.cpu().numpy())\n",
    "            all_labels.extend(labels.cpu().numpy())\n",
    "            all_outputs.extend(probs.cpu().numpy())\n",
    "\n",
    "    all_preds = np.array(all_preds)\n",
    "    all_labels = np.array(all_labels)\n",
    "    all_outputs = np.array(all_outputs)\n",
    "\n",
    "    # Calculate base metrics\n",
    "    accuracy = test_correct / test_total\n",
    "    precision = precision_score(all_labels, all_preds)\n",
    "    recall = recall_score(all_labels, all_preds)\n",
    "    f1 = f1_score(all_labels, all_preds)\n",
    "\n",
    "    tn, fp, fn, tp = confusion_matrix(all_labels, all_preds).ravel()\n",
    "    specificity = tn / (tn + fp)\n",
    "    npv = tn / (tn + fn)\n",
    "\n",
    "    # Calculate AUC and AP\n",
    "    fpr, tpr, _ = roc_curve(all_labels, all_outputs)\n",
    "    roc_auc = auc(fpr, tpr)\n",
    "    precision_curve, recall_curve, _ = precision_recall_curve(all_labels, all_outputs)\n",
    "    pr_auc = average_precision_score(all_labels, all_outputs)\n",
    "\n",
    "    # Generate classification report\n",
    "    class_report = classification_report(all_labels, all_preds, target_names=['Negative', 'Positive'])\n",
    "\n",
    "\n",
    "    # Save results\n",
    "    if save_results:\n",
    "        result_file = os.path.join(results_dir, f\"evaluation_results_{timestamp}.txt\")\n",
    "        with open(result_file, 'w', encoding='utf-8') as f:\n",
    "            # Write basic information\n",
    "            f.write(\"=\" * 50 + \"\\n\")\n",
    "            f.write(\"Pneumonia X-ray Image Classification Model Evaluation Results\\n\")\n",
    "            f.write(f\"Evaluation Time: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\\n\")\n",
    "            f.write(\"=\" * 50 + \"\\n\\n\")\n",
    "\n",
    "            # Write performance metrics\n",
    "            f.write(\"Basic Performance Metrics:\\n\")\n",
    "            f.write(\"-\" * 40 + \"\\n\")\n",
    "            f.write(f\"Test Loss: {test_loss/test_total:.4f}\\n\")\n",
    "            f.write(f\"Accuracy (Accuracy): {accuracy*100:.2f}%\\n\")\n",
    "            f.write(f\"Precision (Precision): {precision:.4f}\\n\")\n",
    "            f.write(f\"Recall/Sensitivity (Recall/Sensitivity): {recall:.4f}\\n\")\n",
    "            f.write(f\"Specificity (Specificity): {specificity:.4f}\\n\")\n",
    "            f.write(f\"F1 Score: {f1:.4f}\\n\")\n",
    "            f.write(f\"Negative Predictive Value (NPV): {npv:.4f}\\n\")\n",
    "            f.write(f\"AUC: {roc_auc:.4f}\\n\")\n",
    "            f.write(f\"PR-AUC: {pr_auc:.4f}\\n\\n\")\n",
    "\n",
    "            # Write confusion matrix\n",
    "            f.write(\"Confusion Matrix:\\n\")\n",
    "            f.write(\"-\" * 40 + \"\\n\")\n",
    "            f.write(\"Prediction\\\\True  Negative(0)  Positive(1)\\n\")\n",
    "            f.write(f\"Negative(0)    {cm[0, 0]}      {cm[0, 1]}\\n\")\n",
    "            f.write(f\"Positive(1)    {cm[1, 0]}      {cm[1, 1]}\\n\\n\")\n",
    "\n",
    "            # Write class detailed performance\n",
    "            f.write(\"Class Detailed Performance:\\n\")\n",
    "            f.write(\"-\" * 40 + \"\\n\")\n",
    "            f.write(f\"Accuracy of Negative Samples: {(cm[0, 0]/(cm[0, 0]+cm[1, 0]))*100:.2f}% ({cm[0, 0]}/{cm[0, 0]+cm[1, 0]})\\n\")\n",
    "            f.write(f\"Accuracy of Positive Samples: {(cm[1, 1]/(cm[1, 1]+cm[0, 1]))*100:.2f}% ({cm[1, 1]}/{cm[1, 1]+cm[0, 1]})\\n\\n\")\n",
    "\n",
    "            # Write classification report\n",
    "            f.write(\"Classification Report:\\n\")\n",
    "            f.write(\"-\" * 40 + \"\\n\")\n",
    "            f.write(class_report + \"\\n\\n\")\n",
    "\n",
    "            # Write evaluation conclusion\n",
    "            f.write(\"Evaluation Conclusion:\\n\")\n",
    "            f.write(\"-\" * 40 + \"\\n\")\n",
    "            f.write(f\"The model's performance in pneumonia detection is {'excellent' if accuracy > 0.9 else 'good' if accuracy > 0.8 else 'fair'}.\\n\")\n",
    "            f.write(f\"It is worth noting that the false negative rate is {fn/(fn+tp):.4f}, which is particularly important in medical diagnosis.\\n\")\n",
    "\n",
    "            # Record generated chart files\n",
    "            f.write(\"\\nGenerated Chart Files:\\n\")\n",
    "            f.write(\"-\" * 40 + \"\\n\")\n",
    "\n",
    "    # Console output\n",
    "    print(f\"\\nTest Results:\")\n",
    "    print(f\"  Loss: {test_loss/test_total:.4f}\")\n",
    "    print(f\"  Accuracy: {accuracy*100:.2f}%\")\n",
    "    print(f\"  Precision: {precision:.4f}\")\n",
    "    print(f\"  Recall: {recall:.4f}\")\n",
    "    print(f\"  Specificity: {specificity:.4f}\")\n",
    "    print(f\"  F1 Score: {f1:.4f}\")\n",
    "    print(f\"  AUC: {roc_auc:.4f}\")\n",
    "\n",
    "    # Confusion matrix display\n",
    "    print(\"\\nConfusion Matrix:\")\n",
    "    print(\"Prediction\\\\True  Negative(0)  Positive(1)\")\n",
    "    print(f\"Negative(0)    {cm[0, 0]}      {cm[0, 1]}\")\n",
    "    print(f\"Positive(1)    {cm[1, 0]}      {cm[1, 1]}\")\n",
    "\n",
    "    # Class detailed performance\n",
    "    neg_acc = cm[0, 0]/(cm[0, 0]+cm[1, 0])\n",
    "    pos_acc = cm[1, 1]/(cm[1, 1]+cm[0, 1])\n",
    "    print(f\"\\n  Accuracy of Negative Samples: {neg_acc*100:.2f}% ({cm[0, 0]}/{cm[0, 0]+cm[1, 0]})\")\n",
    "    print(f\"  Accuracy of Positive Samples: {pos_acc*100:.2f}% ({cm[1, 1]}/{cm[1, 1]+cm[0, 1]})\")\n",
    "\n",
    "\n",
    "    # Return to all assessment indicators\n",
    "    evaluation_results = {\n",
    "        'accuracy': accuracy,\n",
    "        'precision': precision,\n",
    "        'recall': recall,\n",
    "        'specificity': specificity,\n",
    "        'f1': f1,\n",
    "        'npv': npv,\n",
    "        'auc': roc_auc,\n",
    "        'pr_auc': pr_auc,\n",
    "        'confusion_matrix': cm,\n",
    "        'loss': test_loss/test_total,\n",
    "        'class_report': class_report,\n",
    "        'fpr': fpr,\n",
    "        'tpr': tpr,\n",
    "        'pr_curve': precision_curve,\n",
    "        'recall_curve': recall_curve,\n",
    "        'all_preds':all_preds,\n",
    "        'all_labels':all_labels,\n",
    "        'all_outputs':all_outputs\n",
    "    }\n",
    "\n",
    "    return evaluation_results"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "c8f8514957c708de",
   "metadata": {},
   "source": "## 12.Run"
  },
  {
   "cell_type": "code",
   "id": "4b4767e957c7a84d",
   "metadata": {},
   "source": [
    "# Load Data\n",
    "print(\"Load Data...\")\n",
    "train_loader, val_loader, test_loader = load_and_preprocess_data(\n",
    "    data_path, batch_size=64\n",
    ")\n",
    "model = AlterNet_LC(\n",
    "        in_channels=3,\n",
    "        num_classes=2,\n",
    "        blocks_per_stage=[3, 4, 6, 3],\n",
    "        initial_filters=128,\n",
    "        head_counts=[6, 12, 24],\n",
    "        head_dim=64,\n",
    "        window_size=7,\n",
    "        dropout_rate=0.3\n",
    "    ).to(device)\n",
    "\n",
    "# Print model information\n",
    "total_params = sum(p.numel() for p in model.parameters())\n",
    "print(f\"Total model parameters: {total_params:,}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "d8b1d64314abacbf",
   "metadata": {},
   "source": [
    "# Train model\n",
    "print(\"Start Training...\")\n",
    "model, history = train_model(\n",
    "    model,\n",
    "    train_loader,\n",
    "    val_loader,\n",
    "    epochs=100,\n",
    "    lr=5e-5,\n",
    "    weight_decay=1.6e-4\n",
    ")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "fb7b5b73d4b92756",
   "metadata": {},
   "source": [
    "plot_training(history)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "b4ff7d18181b8105",
   "metadata": {},
   "source": [
    "evaluate_results = evaluate_model(model, test_loader)\n",
    "fpr = evaluate_results['fpr']\n",
    "tpr = evaluate_results['tpr']\n",
    "pr_auc = evaluate_results['pr_auc']\n",
    "roc_auc = auc(fpr, tpr)\n",
    "recall_curve = evaluate_results['recall_curve']\n",
    "precision_curve = evaluate_results['pr_curve']\n",
    "all_labels = evaluate_results['all_labels']"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "23afdb938ce6174d",
   "metadata": {},
   "source": [
    "# 1. Draw ROC curve\n",
    "plt.figure(figsize=(10, 8))\n",
    "plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC Curve (AUC = {roc_auc:.4f})')\n",
    "plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--', label='random guess')\n",
    "plt.xlim([0.0, 1.0])\n",
    "plt.ylim([0.0, 1.05])\n",
    "plt.xlabel('false positive rate')\n",
    "plt.ylabel('true positive rate')\n",
    "plt.title('ROC Curve')\n",
    "plt.legend(loc=\"lower right\")\n",
    "plt.grid(True, linestyle='--', alpha=0.7)\n",
    "\n",
    "plt.show()"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "# 2. Draw PR curve\n",
    "plt.figure(figsize=(10, 8))\n",
    "plt.plot(recall_curve, precision_curve, color='green', lw=2,\n",
    "            label=f'PR Curve (AP = {pr_auc:.4f})')\n",
    "plt.axhline(y=np.sum(all_labels) / len(all_labels), color='navy',\n",
    "            linestyle='--', label='random guess')\n",
    "plt.xlim([0.0, 1.0])\n",
    "plt.ylim([0.0, 1.05])\n",
    "plt.xlabel('recall rate')\n",
    "plt.ylabel('precision rate')\n",
    "plt.title('PR Curve')\n",
    "plt.legend(loc=\"lower left\")\n",
    "plt.grid(True, linestyle='--', alpha=0.7)\n",
    "\n",
    "plt.show()"
   ],
   "id": "7ac98a1b1035d698"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "# Saving the final model\n",
    "print(\"Saving the final model...\")\n",
    "torch.save(model.state_dict(), f'alternet_lc_{datetime.datetime.now().strftime(\"%Y%m%d\")}.pth')\n",
    "print(f\"Final model saved: alternet_lc_{datetime.datetime.now().strftime('%Y%m%d')}.pth\")"
   ],
   "id": "91acade911f470f9"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "source": "",
   "id": "448fb3d9a2cf57ca",
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
